package com.xpalive.utils;

import com.xpalive.exception.BizException;
import java.util.List;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

public class ThreadPoolCountDown {

    private final ExecutorService executorService;
    private final static int cpuCount = Runtime.getRuntime().availableProcessors();

    public ThreadPoolCountDown() {
        this(cpuCount, cpuCount * 2, 120L, TimeUnit.MINUTES,
            new ArrayBlockingQueue<>(1000));
    }

    public ThreadPoolCountDown(int corePoolSize,
        int maximumPoolSize,
        long keepAliveTime,
        TimeUnit unit,
        BlockingQueue<Runnable> workQueue) {
        this.executorService = new ThreadPoolExecutor(corePoolSize,
            maximumPoolSize, keepAliveTime, unit, workQueue);
    }

    private CountDownLatch countDownLatch;

    public Executor tasks(List<? extends TaskCountDownLatch> tasks) {
        countDownLatch = new CountDownLatch(tasks.size());
        return new Executor(tasks);
    }


    public class Executor {

        private List<? extends TaskCountDownLatch> tasks;

        private Executor(List<? extends TaskCountDownLatch> tasks) {
            this.tasks = tasks;
        }

        public void exec() {
            for (TaskCountDownLatch task : tasks) {
                task.setCountDownLatch(countDownLatch);
                executorService.execute(task);
            }
            try {
                countDownLatch.await();
            } catch (InterruptedException e) {
                throw BizException.of(e);
            }
        }
    }

    public static abstract class TaskCountDownLatch implements Runnable {

        private CountDownLatch countDownLatch;

        private void setCountDownLatch(CountDownLatch countDownLatch) {
            this.countDownLatch = countDownLatch;
        }

        @Override
        public void run() {
            try {
                run0();
            } finally {
                countDownLatch.countDown();
            }
        }

        public abstract void run0();
    }
}
